{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Appendix 10.2: Tool Use\n",
    "\n",
    "- [Lesson](#lesson)\n",
    "- [Exercises](#exercises)\n",
    "\n",
    "## Setup\n",
    "\n",
    "Run the following setup cell to load your API key and establish the `get_completion` helper function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import python's built-in regular expression library\n",
    "import re\n",
    "from ollama import chat, ChatResponse, Options\n",
    "\n",
    "\n",
    "# Stores the MODEL_NAME variable from the IPython store\n",
    "%store -r MODEL_NAME\n",
    "\n",
    "def get_completion(prompt: str, system_prompt=\"\", prefill=\"\", tools=None):\n",
    "    # Ensure tools is a list, even if None is passed\n",
    "    tools = tools if tools is not None else []\n",
    "    \n",
    "    response = chat(\n",
    "        model=MODEL_NAME,\n",
    "        options=Options(\n",
    "            max_tokens=2000,\n",
    "            temperature=0.0,\n",
    "            num_ctx=8192 # larger context defined to include large prompt in this chapter, default value is 2048\n",
    "        ),\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},  \n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "            {\"role\": \"assistant\", \"content\": prefill}\n",
    "        ],\n",
    "        tools=tools,\n",
    "    )\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Lesson\n",
    "\n",
    "While it might seem conceptually complex at first, tool use, a.k.a. function calling, is actually quite simple! You already know all the skills necessary to implement tool use, which is really just a combination of substitution and prompt chaining.\n",
    "\n",
    "In previous substitution exercises, we substituted text into prompts. With tool use, we substitute tool or function results into prompts. Ollama can't literally call or access tools and functions. Instead, we have Ollama:\n",
    "1. Output the tool name and arguments it wants to call\n",
    "2. Halt any further response generation while the tool is called\n",
    "3. Then we reprompt with the appended tool results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Function calling is useful because it expands Ollama's capabilities and enables Ollama to handle much more complex, multi-step tasks.\n",
    "Some examples of functions you can give Ollama:\n",
    "- Calculator\n",
    "- Word counter\n",
    "- SQL database querying and data retrieval\n",
    "- Weather API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can get Ollama to do tool use in two way:\n",
    "- Manual\n",
    "- Automatic\n",
    "  \n",
    "Python functions or any other functions you want to call can be passed to the model through the tools parameter of the chat API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examples\n",
    "\n",
    "To enable tool use in Ollama, we start with defining the functions we are going to use. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_two_numbers(a: int, b: int) -> int:\n",
    "  \"\"\"\n",
    "  Add two numbers\n",
    "\n",
    "  Args:\n",
    "    a: The first integer number\n",
    "    b: The second integer number\n",
    "\n",
    "  Returns:\n",
    "    int: The sum of the two numbers\n",
    "  \"\"\"\n",
    "  return a + b\n",
    "\n",
    "def subtract_two_numbers(a: int, b: int) -> int:\n",
    "  \"\"\"\n",
    "  Subtract two numbers\n",
    "  \"\"\"\n",
    "  return a - b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we add them to a dictionary to be able to call them when Model will ask us to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "available_functions = {\n",
    "  'add_two_numbers': add_two_numbers,\n",
    "  'subtract_two_numbers': subtract_two_numbers,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here the manual creation of JSON schema for the subtract_two_numbers function, while for add_two_numbers we will use automatic way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tools can still be manually defined and passed into chat\n",
    "subtract_two_numbers_tool = {\n",
    "  'type': 'function',\n",
    "  'function': {\n",
    "    'name': 'subtract_two_numbers',\n",
    "    'description': 'Subtract two numbers',\n",
    "    'parameters': {\n",
    "      'type': 'object',\n",
    "      'required': ['a', 'b'],\n",
    "      'properties': {\n",
    "        'a': {'type': 'integer', 'description': 'The first number'},\n",
    "        'b': {'type': 'integer', 'description': 'The second number'},\n",
    "      },\n",
    "    },\n",
    "  },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, use the *tools* field to pass the functions as tools to Ollama and we use the returned tool call and arguments provided by the model to call the respective function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prompt\n",
    "PROMPT = \"What is 2015 + 20?\"\n",
    "# Get Ollama's response\n",
    "response = get_completion(PROMPT, tools=[add_two_numbers,subtract_two_numbers])\n",
    "\n",
    "for tool in response.message.tool_calls or []:\n",
    "  function_to_call = available_functions.get(tool.function.name)\n",
    "  if function_to_call:\n",
    "    result = function_to_call(**tool.function.arguments)\n",
    "  else:\n",
    "    print('Function not found:', tool.function.name)\n",
    "\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's chain the result for the first call with another one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prompt\n",
    "PROMPT = f\"{result} - 10?\"\n",
    "\n",
    "# Get Ollama's response\n",
    "response = get_completion(PROMPT, tools=[add_two_numbers,subtract_two_numbers])\n",
    "\n",
    "for tool in response.message.tool_calls or []:\n",
    "  function_to_call = available_functions.get(tool.function.name)\n",
    "  if function_to_call:\n",
    "    result = function_to_call(**tool.function.arguments)\n",
    "  else:\n",
    "    print('Function not found:', tool.function.name)\n",
    "\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pass existing functions as tools\n",
    "\n",
    "Functions from existing Python libraries, SDKs, and elsewhere can now also be provided as tools. For example, the following code passes the request function from the requests library as a tool to fetch the contents of the Ollama website:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "available_functions = {\n",
    "  'request': requests.request,\n",
    "}\n",
    "\n",
    "PROMPT = \"get the ollama.com webpage?\"\n",
    "\n",
    "# Print Ollama's response\n",
    "response = get_completion(PROMPT, tools=[requests.request])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for tool in response.message.tool_calls or []:\n",
    "  function_to_call = available_functions.get(tool.function.name)\n",
    "  if function_to_call == requests.request:\n",
    "    # Make an HTTP request to the URL specified in the tool call\n",
    "    resp = function_to_call(\n",
    "      method=tool.function.arguments.get('method'),\n",
    "      url=tool.function.arguments.get('url'),\n",
    "    )\n",
    "    print(resp.text)\n",
    "  else:\n",
    "    print('Function not found:', tool.function.name)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Exercises\n",
    "- [Exercise 10.2.1 - SQL](#exercise-1021---SQL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10.2.1 - SQL\n",
    "In this exercise, you'll be writing a tool use prompt for querying and writing to the world's smallest \"database\". Here's the initialized database, which is really just a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = {\n",
    "    \"users\": [\n",
    "        {\"id\": 1, \"name\": \"Alice\", \"email\": \"alice@example.com\"},\n",
    "        {\"id\": 2, \"name\": \"Bob\", \"email\": \"bob@example.com\"},\n",
    "        {\"id\": 3, \"name\": \"Charlie\", \"email\": \"charlie@example.com\"}\n",
    "    ],\n",
    "    \"products\": [\n",
    "        {\"id\": 1, \"name\": \"Widget\", \"price\": 9.99},\n",
    "        {\"id\": 2, \"name\": \"Gadget\", \"price\": 14.99},\n",
    "        {\"id\": 3, \"name\": \"Doohickey\", \"price\": 19.99}\n",
    "    ]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here is the code for the functions that write to and from the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_user(user_id):\n",
    "    for user in db[\"users\"]:\n",
    "        if user[\"id\"] == user_id:\n",
    "            return user\n",
    "    return None\n",
    "\n",
    "def get_product(product_id):\n",
    "    for product in db[\"products\"]:\n",
    "        if product[\"id\"] == product_id:\n",
    "            return product\n",
    "    return None\n",
    "\n",
    "def add_user(name, email):\n",
    "    user_id = len(db[\"users\"]) + 1\n",
    "    user = {\"id\": user_id, \"name\": name, \"email\": email}\n",
    "    db[\"users\"].append(user)\n",
    "    return user\n",
    "\n",
    "def add_product(name, price):\n",
    "    product_id = len(db[\"products\"]) + 1\n",
    "    product = {\"id\": product_id, \"name\": name, \"price\": price}\n",
    "    db[\"products\"].append(product)\n",
    "    return product"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To solve the exercise, start by defining a system prompt like `system_prompt_tools_specific_tools` above. Make sure to include the name and description of each tool, along with the name and type and description of each parameter for each function. We've given you some starting scaffolding below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you're ready, you can try out your tool definition system prompt on the examples below. Just run the below cell!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = [\n",
    "    \"Add a user to the database named Debora with email debora@example.com.\",\n",
    "    \"Add a product to the database named Thingo with price 10\",\n",
    "    \"Tell me the name of User 2 from the database\",\n",
    "    \"Tell me the name of Product 3 from the database\"\n",
    "]\n",
    "\n",
    "available_functions = {\n",
    "    'get_user': get_user,\n",
    "    'get_product': get_product,\n",
    "    'add_user': add_user,\n",
    "    'add_product': add_product,\n",
    "}\n",
    "\n",
    "for example in examples:\n",
    "    # Get & print Ollama's response\n",
    "    function_calling_response = get_completion(example, tools=[get_user, get_product, add_user, add_product])\n",
    "    print(example, \"\\n----------\\n\\n\", function_calling_response, \"\\n*********\\n*********\\n*********\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you did it right, the function calling messages should call the `add_user`, `add_product`, `get_user`, and `get_product` functions correctly.\n",
    "\n",
    "For extra credit, add some code cells and write parameter-parsing code. Then call the functions with the parameters Ollama gives you to see the state of the \"database\" after the call."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "❓ If you want to see a possible solution, run the cell below!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hints import exercise_10_2_1_solution; print(exercise_10_2_1_solution)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Congrats!\n",
    "\n",
    "Congratulations on learning tool use and function calling! \n",
    "For more information on how to leverage function calling go to blog post **[Ollama Python library 0.4 with function calling improvements](https://ollama.com/blog/functions-as-tools)**"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
